from numpy import *

# 创建一个用于测试的简单数据集
def loadDataSet():
    return [[1, 3, 4], [2, 3, 5], [1, 2, 3, 5], [2, 5]]

# 构建第一个候选集的列表C1
def createC1(dataSet):
    C1 = []                             # 创建一个空列表C1，用来存储所有不重复的项值
    for transaction in dataSet:         # 对数据集中的每条记录做扫描
        for item in transaction:        # 对记录中的每个项值做扫描
            if not [item] in C1:        # 如果该项值的集合不存在于列表C1中，
                C1.append([item])       # 则将该项值的集合添加到列表C1
                
    C1.sort()                           # 对列表C1进行排序
    return map(frozenset, C1)           # 将列表C1设置为冰冻列表（不可修改），
                                        # 其每个单元作为字典键值使用
                            
# 三个输入参数，数据集D、候选项集列表Ck、感兴趣项集的最小支持度minSupport
def scanD(D, Ck, minSupport):
    # 扫描数据集
    ssCnt = {}                      # 创建一个空字典ssCnt
    for tid in D:                   # 对数据集中的每条记录做扫描
        for can in Ck:              # 对候选集中的每个集合做扫描
            if can.issubset(tid):   # 如果候选集中的集合是记录中的一部分，那么增加字典中对应的计数值
                if not can in ssCnt: ssCnt[can]=1
                else: ssCnt[can] += 1
    # 计算支持度
    numItems = float(len(D))        # 获取数据集中记录的条数numItems
    retList = []                    # 构建一个空列表retList，该列表包含满足最小支持度要求的集合
    supportData = {}                # 构建一个空字典supportData，用于存储最频繁项集的支持度
    for key in ssCnt:                       # 对遍历字典ssCnt中的每一个元素
        support = ssCnt[key]/numItems       # 计算每个候选集的支持度
        if support >= minSupport:           # 如果候选集的支持度大于minSupport
            retList.insert(0,key)           # 将该候选集插入retList列表中
        supportData[key] = support          # 保存每个候选集的支持度
    return retList, supportData         # 返回满足最小支持度要求的候选集集合列表retList
                                        # 和一个包含所有候选集支持度值的字典supportData
                                        
# 创建候选集列表Ck，输入参数，频繁项集列表Lk和项集元素个数k
def aprioriGen(Lk, k): 
    retList = []        # 构建一个空列表retList
    lenLk = len(Lk)     # 获取频繁项集列表Lk的元素个数
    for i in range(lenLk):  
        for j in range(i+1, lenLk): 
            L1 = list(Lk[i])[:k-2]; L2 = list(Lk[j])[:k-2]
            L1.sort(); L2.sort()
            if L1==L2: # 如果两个集合中的前k-2个元素值都相等
                retList.append(Lk[i] | Lk[j]) # 将这两个集合取并集，并存入列表retList中
    return retList      # 返回候选集列表Ck

# Apriori算法
def apriori(dataSet, minSupport = 0.5):
    C1 = list(createC1(dataSet))        # 构建第一个候选集的列表C1
    D = list(map(set, dataSet))         # 构建集合表示的数据集D
    L1, supportData = scanD(D, C1, minSupport)  # 构建频繁项集的列表L1
    L = [L1]        # 构造列表L，用于存放L1、L2、L3...
    k = 2           # 初始化K值
    while (len(L[k-2]) > 0):        # 当集合中项的个数大于0时
        Ck = aprioriGen(L[k-2], k)              # 创建候选集列表Ck
        Lk, supK = scanD(D, Ck, minSupport)     # 扫描数据集D，构建频繁项集的列表Lk，并计算支持度
        supportData.update(supK)                # 更新支持度数据
        L.append(Lk)            # 将Lk存入列表L
        k += 1
    return L, supportData

# 关联规则生成函数，
# 三个输入参数，频繁项集列表L、包含哪些频繁项集支持数据的字典supportData和最小可信度阈值minConf
def generateRules(L, supportData, minConf=0.7):  # 列表L和字典supportData都来自函数apriori()
    bigRuleList = []    # 构造一个空列表bigRuleList，用于存放规则
    for i in range(1, len(L)):  # 只获取有两个或更多元素的集合
        for freqSet in L[i]:    # 对集合中的每个元素进行扫描
            H1 = [frozenset([item]) for item in freqSet]
            if (i > 1): # 如果频繁项集的元素数目超过2，调用rulesFromConseq()函数做进一步合并
                rulesFromConseq(freqSet, H1, supportData, bigRuleList, minConf)
            else:       # 如果频繁项集只有两个元素，则调用calcConf()函数来计算可信度值
                calcConf(freqSet, H1, supportData, bigRuleList, minConf)
    return bigRuleList         

# 计算规则的可信度
def calcConf(freqSet, H, supportData, brl, minConf=0.7):
    prunedH = [] # 创建一个空列表prunedH
    for conseq in H:    # 遍历H中的所有项集
        conf = supportData[freqSet]/supportData[freqSet-conseq] # 计算可信度
        if conf >= minConf:  # 如果可信度大于最小可信度，则创建规则，并将其保存到brl列表中
            print (freqSet-conseq,'-->',conseq,'conf:',conf)
            brl.append((freqSet-conseq, conseq, conf))
            prunedH.append(conseq)
    return prunedH

# 生成候选规则函数
def rulesFromConseq(freqSet, H, supportData, brl, minConf=0.7):
    m = len(H[0])
    if (len(freqSet) > (m + 1)): # 检查频繁项集是否大到可以移除大小为m的子集
        Hmp1 = aprioriGen(H, m+1) # 函数aprioriGen()被用来生成H中元素的无重复组合
        Hmp1 = calcConf(freqSet, Hmp1, supportData, brl, minConf)
        if (len(Hmp1) > 1):    #need at least two sets to merge
            rulesFromConseq(freqSet, Hmp1, supportData, brl, minConf)
            
